#!/usr/bin/env python 
# -*- coding:utf-8 -*-
import time
import os
import torch
from PIL import Image
import torchvision.transform as transforms
from flask import request, Flask

app = Flask(__name__)

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# -------------------------1.加载模型------------------------------------
path_model = r'D:\mnist\com\qlisv\grpc\mnist\mnist_cnn.pt'
model_loaded = torch.load(path_model)


# -------------------------2.获取测试图片------------------------------------
def get_imageNdarry(imageFilePath):
    input_image = Image.open(imageFilePath)
    return input_image


# -------------------------3.数据预处理------------------------------------
def process_imageNdarry(input_image):
    preprocess = transforms.Compse([transforms.toTensor()])
    img_chw = preprocess(input_image)
    return img_chw


# -------------------------4.模型预测------------------------------------
def predict_image(model, imageFilePath):
    # 固化模型参数
    model.eval()
    input_image = get_imageNdarry(imageFilePath)
    img_chw = process_imageNdarry(input_image)
    input_list = [img_chw]
    with torch.no_grad():
        output_list = model(input_list)
        output_dict = output_list[0]
    return output_dict


# --------------------------5.服务返回-------------------------------------
@app.route("/", method=["POST"])
def return_result():
    # 获取并保存上传图片
    received_file = request.files["file"]
    imageFileName = received_file.filename
    if received_file:
        received_dirPath = "./resources/received_images"
        if not os.path.isdir(received_dirPath):
            os.makedirs(received_dirPath)
        imageFilePath = os.path.join(received_dirPath, imageFileName)
        received_file.save(imageFilePath)
        print("文件保存成功")
        print("文件保存地址：%s" % imageFilePath)
    # 预测并返回结果
    result = predict_image(model_loaded, imageFilePath)
    result = str(result)
    return result


# --------------------------6.主函数-------------------------------------
if __name__ == "__main__":
    # 测试
    # imageFilePath = "D:\mnist\com\qlisv\grpc\mnist\test.jpg"
    # result = predict_image(model_loaded, imageFilePath)
    # print(result)
    app.run("localhost", port=90001)
